{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "fc539b17",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef70f299",
   "metadata": {},
   "source": [
    "线性模型表示了一个超平面。模型的某次特定预测可以表示为 $\\hat y = \\bold {x}^{\\top} \\bold{w} + b$，而一堆预测可以表示为类似 $\\hat{\\bold{y}} = \\bold X \\bold w + b$ 的形式。当然，不需要管太多细节。\n",
    "\n",
    "深度学习的任务通常并不是最小化 loss，最大化 likelihood，更艰巨的任务在于泛化性能。\n",
    "\n",
    "所以不必对每次 batch_size = minibatchlen 而不是 train_data.full_len 这件事太过芥蒂，毕竟我们要解决的更多是无解析解的问题，追求无偏的 loss 空间意义不大。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "8e1db335",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SyntheticRegressionData():  #@save\n",
    "    \"\"\"Synthetic data for linear regression.\"\"\"\n",
    "    def __init__(self, w, b, noise=0.01, num_train=1000, num_val=1000,\n",
    "                 batch_size=32):\n",
    "        super().__init__()\n",
    "        device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else \"cpu\"\n",
    "        n = num_train + num_val # 这里完全没有必要，但从 d2l 里薅的代码就不要求这么多了\n",
    "        self.X = torch.randn(n, len(w))  # 生成 n 个样本，每个样本 len(w) 维\n",
    "        noise = torch.randn(n, 1) * noise\n",
    "        self.y = torch.matmul(self.X, w.reshape((-1, 1))) + b + noise\n",
    "        self.y = self.y.to(device)\n",
    "        self.X = self.X.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "b4bec9e6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "features: tensor([-0.6742,  0.0689,  0.8240], device='cuda:0') \n",
      "label: tensor([84.1889], device='cuda:0')\n",
      "2000 2000\n"
     ]
    }
   ],
   "source": [
    "data = SyntheticRegressionData(w=torch.tensor([2, -3.4, 99]), b=4.2)\n",
    "print('features:', data.X[0],'\\nlabel:', data.y[0])\n",
    "print(len(data.X), len(data.y))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f619dfc8",
   "metadata": {},
   "source": [
    "要搞 torch 框架，就得熟悉它的 DataSet 类，写一个。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "3eff9d18",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyDataset(torch.utils.data.Dataset): # py 的继承这么写\n",
    "    def __init__(self):\n",
    "        self.data = SyntheticRegressionData(w=torch.tensor([2, -3.4, 99]), b=4.2)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data.X)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.data.X[idx], self.data.y[idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "c7e56d9f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Feature batch shape: torch.Size([32, 3])\n",
      "Labels batch shape: torch.Size([32, 1])\n"
     ]
    }
   ],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "train_dataloader = DataLoader(MyDataset(), batch_size=32, shuffle=True)\n",
    "test_dataloader = DataLoader(MyDataset(), batch_size=32, shuffle=True)\n",
    "\n",
    "train_features, train_labels = next(iter(train_dataloader))\n",
    "print(f\"Feature batch shape: {train_features.size()}\")\n",
    "print(f\"Labels batch shape: {train_labels.size()}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "d48cbb4e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using cuda device\n"
     ]
    }
   ],
   "source": [
    "from torch import nn\n",
    "device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else \"cpu\"\n",
    "print(f\"Using {device} device\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4caa3005",
   "metadata": {},
   "source": [
    "官方文档写了子类写法。自定义网络显然还是更适合学习的，毕竟看接口太烦了。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "5bb815d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyNN(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.linear = nn.Linear(3, 1)\n",
    "        # 不需要 Flatten，不需要 Sequential，朴实无华\n",
    "\n",
    "    def forward(self, x): # 推理\n",
    "        return self.linear(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "9a04ed6e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MyNN(\n",
      "  (linear): Linear(in_features=3, out_features=1, bias=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "model = MyNN().to(device)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "2af352ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_rate = 1e-3\n",
    "batch_size = 64\n",
    "epochs = 5\n",
    "loss_fn = nn.MSELoss(reduction='mean')\n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "4ada2279",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_loop(dataloader, model, loss_fn, optimizer):\n",
    "    size = len(dataloader.dataset)\n",
    "    # Set the model to training mode - important for batch normalization and dropout layers\n",
    "    # Unnecessary in this situation but added for best practices\n",
    "    model.train()\n",
    "    for batch, (X, y) in enumerate(dataloader):\n",
    "        # Compute prediction and loss\n",
    "        pred = model(X)\n",
    "        loss = loss_fn(pred, y)\n",
    "\n",
    "        # Backpropagation\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        if batch % 100 == 0:\n",
    "            loss, current = loss.item(), batch * batch_size + len(X)\n",
    "            print(f\"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]\")\n",
    "\n",
    "\n",
    "def test_loop(dataloader, model, loss_fn):\n",
    "    # Set the model to evaluation mode - important for batch normalization and dropout layers\n",
    "    # Unnecessary in this situation but added for best practices\n",
    "    model.eval()\n",
    "    size = len(dataloader.dataset)\n",
    "    num_batches = len(dataloader)\n",
    "    test_loss, correct = 0, 0\n",
    "\n",
    "    # Evaluating the model with torch.no_grad() ensures that no gradients are computed during test mode\n",
    "    # also serves to reduce unnecessary gradient computations and memory usage for tensors with requires_grad=True\n",
    "    with torch.no_grad():\n",
    "        for X, y in dataloader:\n",
    "            pred = model(X)\n",
    "            test_loss += loss_fn(pred, y).item()\n",
    "            correct += (pred.argmax(1) == y).type(torch.float).sum().item()\n",
    "\n",
    "    test_loss /= num_batches\n",
    "    correct /= size\n",
    "    print(f\"Test Error: \\n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "0a1758c6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1\n",
      "-------------------------------\n",
      "loss: 0.034477  [   32/ 2000]\n",
      "Test Error: \n",
      " Accuracy: 0.0%, Avg loss: 0.032732 \n",
      "\n",
      "Epoch 2\n",
      "-------------------------------\n",
      "loss: 0.067478  [   32/ 2000]\n",
      "Test Error: \n",
      " Accuracy: 0.0%, Avg loss: 0.025759 \n",
      "\n",
      "Epoch 3\n",
      "-------------------------------\n",
      "loss: 0.029091  [   32/ 2000]\n",
      "Test Error: \n",
      " Accuracy: 0.0%, Avg loss: 0.020031 \n",
      "\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "epochs = 3\n",
    "for t in range(epochs):\n",
    "    print(f\"Epoch {t+1}\\n-------------------------------\")\n",
    "    train_loop(train_dataloader, model, loss_fn, optimizer)\n",
    "    test_loop(test_dataloader, model, loss_fn)\n",
    "print(\"Done!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "c8c7909c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parameter containing:\n",
      "tensor([[ 2.0276, -3.3910, 98.8595]], device='cuda:0', requires_grad=True) Parameter containing:\n",
      "tensor([4.1960], device='cuda:0', requires_grad=True)\n"
     ]
    }
   ],
   "source": [
    "print(model.linear.weight, model.linear.bias)\n",
    "# 作为参考，当时是以 w=torch.tensor([2, -3.4, 99]), b=4.2 带上一些正态分布噪声生成的训练数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "711d35fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), 'model_weights.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "e8ee01e1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parameter containing:\n",
      "tensor([[ 2.0276, -3.3910, 98.8595]], device='cuda:0', requires_grad=True) Parameter containing:\n",
      "tensor([4.1960], device='cuda:0', requires_grad=True)\n"
     ]
    }
   ],
   "source": [
    "new_model = MyNN().to(device)\n",
    "new_model.load_state_dict(torch.load('model_weights.pth'))\n",
    "print(new_model.linear.weight, new_model.linear.bias)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
